﻿using System;
using System.Collections.Generic;
using System.Windows.Forms;

namespace ExtensionMethodExtravaganza
{
	public static class TreeViewExtensions
	{
		public static List<TreeNode> GetDescendantNodes(this TreeNode treeNode) 
		{
			return GetDescendantNodes<TreeNode>(treeNode.Nodes);
		}

		public static List<T> GetDescendantNodes<T>(this TreeNode treeNode) where T : TreeNode
		{
			return GetDescendantNodes<T>(treeNode.Nodes);
		}

		public static List<TreeNode> GetDescendantNodes(this TreeView treeView) 
		{
			return GetDescendantNodes<TreeNode>(treeView.Nodes);
		}

		public static List<T> GetDescendantNodes<T>(this TreeView treeView) where T : TreeNode
		{
			return GetDescendantNodes<T>(treeView.Nodes);
		}

		private static List<T> GetDescendantNodes<T>(TreeNodeCollection treeNodes) where T : TreeNode
		{
			var foundNodes = new List<T>();

			GetDescendantNodes(treeNodes, foundNodes);
			
			return foundNodes;
		}

		private static void GetDescendantNodes<T>(TreeNodeCollection treeNodes, ICollection<T> foundNodes) where T : TreeNode
		{
			for (var i = 0; i < treeNodes.Count; i++)
			{
				var node = (T)treeNodes[i];
				foundNodes.Add(node);

				if (node.Nodes.Count > 0)
				{
					GetDescendantNodes(node.Nodes, foundNodes);
				}
			}
		}

		/// <summary>
		/// Finds a node in the TreeView.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static T FindNode<T>(this TreeView treeView, Predicate<T> condition) where T : TreeNode
		{
			return FindNode(treeView.Nodes, condition);
		}

		/// <summary>
		/// Finds a node in the TreeView.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static TreeNode FindNode(this TreeView treeView, Predicate<TreeNode> condition)
		{
			return FindNode<TreeNode>(treeView, condition);
		}

		/// <summary>
		/// Finds nodes in the TreeView.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static List<T> FindNodes<T>(this TreeView treeView, Predicate<T> condition) where T : TreeNode
		{
			var foundNodes = new List<T>();

			FindNodes(treeView.Nodes, condition, foundNodes);

			return foundNodes;
		}

		/// <summary>
		/// Finds nodes in the TreeView.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static List<TreeNode> FindNodes(this TreeView treeView, Predicate<TreeNode> condition)
		{
			return FindNodes<TreeNode>(treeView, condition);
		}

		/// <summary>
		/// Finds a node below the specified TreeNode.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static T FindNode<T>(this TreeNode treeNode, Predicate<T> condition) where T : TreeNode
		{
			return FindNode(treeNode.Nodes, condition);
		}

		/// <summary>
		/// Finds a node below the specified TreeNode.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static TreeNode FindNode(this TreeNode treeNode, Predicate<TreeNode> condition)
		{
			return FindNode<TreeNode>(treeNode, condition);
		}

		/// <summary>
		/// Finds nodes below the specified TreeNode.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static List<T> FindNodes<T>(this TreeNode treeNode, Predicate<T> condition) where T : TreeNode
		{
			var foundNodes = new List<T>();

			FindNodes(treeNode.Nodes, condition, foundNodes);

			return foundNodes;
		}

		/// <summary>
		/// Finds nodes below the specified TreeNode.
		/// </summary>
		/// <remarks>Could be done more efficiently.</remarks>
		public static List<TreeNode> FindNodes(this TreeNode treeNode, Predicate<TreeNode> condition)
		{
			return FindNodes<TreeNode>(treeNode, condition);
		}

		private static T FindNode<T>(TreeNodeCollection nodes, Predicate<T> condition) where T : TreeNode
		{
			T foundNode = null;

			for (var i = 0; i < nodes.Count; i++)
			{
				var node = (T)nodes[i];
				
				foundNode = condition(node) ? node : FindNode(node.Nodes, condition);

				if (foundNode != null)
				{
					break;
				}
			}

			return foundNode;
		}

		private static void FindNodes<T>(TreeNodeCollection nodes, Predicate<T> condition, List<T> foundNodes) where T : TreeNode
		{
			for (var i = 0; i < nodes.Count; i++)
			{
				var node = (T)nodes[i];

				if (condition(node))
				{
					foundNodes.Add(node);
				}

				FindNodes(node.Nodes, condition, foundNodes);
			}
		}
	}
}
